// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions {
namespace {

class MapTopNKeysTest : public test::FunctionBaseTest {};

TEST_F(MapTopNKeysTest, emptyMap) {
  RowVectorPtr input = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>({
          "{}",
      }),
  });

  assertEqualVectors(
      evaluate("map_top_n_keys(c0, 3)", input),
      make_array_vector_from_json<int32_t>({
          "[]",
      }));
}

TEST_F(MapTopNKeysTest, multipleMaps) {
  RowVectorPtr input = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>({
          "{3:1, 2:1, 5:1, 4:1, 1:1}",
          "{3:1, 2:1, 1:1}",
          "{2:1, 1:1}",
      }),
  });

  assertEqualVectors(
      evaluate("map_top_n_keys(c0, 3)", input),
      make_array_vector_from_json<int32_t>({
          "[5, 4, 3]",
          "[3, 2, 1]",
          "[2, 1]",
      }));
}

TEST_F(MapTopNKeysTest, nIsZero) {
  RowVectorPtr input = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>({
          "{2:1, 1:1}",
      }),
  });

  assertEqualVectors(
      evaluate("map_top_n_keys(c0, 0)", input),
      make_array_vector_from_json<int32_t>({"[]"}));
}

TEST_F(MapTopNKeysTest, nIsNegative) {
  RowVectorPtr input = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>({
          "{2:1, 1:1}",
      }),
  });

  POLLUX_ASSERT_THROW(
      evaluate("map_top_n_keys(c0, -1)", input),
      "n must be greater than or equal to 0");
}

TEST_F(MapTopNKeysTest, timestampWithTimeZone) {
  auto testMapTopNKeys = [&](const std::vector<int64_t>& keys,
                             const std::vector<int32_t>& values,
                             const std::vector<int64_t>& expectedKeys) {
    const auto map = make_map_vector(
        {0},
        make_flat_vector(keys, TIMESTAMP_WITH_TIME_ZONE()),
        make_flat_vector(values));
    const auto expected = make_array_vector(
        {0}, make_flat_vector(expectedKeys, TIMESTAMP_WITH_TIME_ZONE()));

    const auto result = evaluate("map_top_n_keys(c0, 3)", make_row_vector({map}));

    assertEqualVectors(expected, result);
  };

  testMapTopNKeys(
      {pack(1, 1), pack(2, 2), pack(3, 3), pack(4, 4), pack(5, 5)},
      {3, 5, 1, 4, 2},
      {pack(5, 5), pack(4, 4), pack(3, 3)});
  testMapTopNKeys(
      {pack(5, 1), pack(4, 2), pack(3, 3), pack(2, 4), pack(1, 5)},
      {3, 5, 1, 4, 2},
      {pack(5, 1), pack(4, 2), pack(3, 3)});
  testMapTopNKeys(
      {pack(3, 1), pack(5, 2), pack(1, 3), pack(4, 4), pack(2, 5)},
      {1, 2, 3, 4, 5},
      {pack(5, 2), pack(4, 4), pack(3, 1)});
  testMapTopNKeys(
      {pack(3, 5), pack(5, 4), pack(4, 2), pack(2, 1)},
      {1, 2, 4, 5},
      {pack(5, 4), pack(4, 2), pack(3, 5)});
  testMapTopNKeys(
      {pack(10, 3), pack(7, 2), pack(0, 1)},
      {1, 2, 3},
      {pack(10, 3), pack(7, 2), pack(0, 1)});
  testMapTopNKeys(
      {pack(1, 10), pack(10, 1)}, {1, 2}, {pack(10, 1), pack(1, 10)});
  testMapTopNKeys({}, {}, {});
}

TEST_F(MapTopNKeysTest, nestedVarcharMap) {
  // Create inner maps: MAP<VARCHAR, VARCHAR>
  auto array_vector = make_flat_vector<std::string>({"k1", "k2", "k3"});
  auto innerMapsVector = make_map_vector_from_json<std::string, std::string>(
      {R"({"a":"test_value", "b":"test_value"})",
       R"({"b":"3", "c":"test_value"})",
       R"({"b":"3", "c":"test_value"})"});

  auto mapvector = make_map_vector({0}, array_vector, innerMapsVector);

  RowVectorPtr input = make_row_vector({mapvector});
  auto result = evaluate("map_top_n_keys(c0, 2)", input);
  auto expected = make_array_vector_from_json<std::string>({"[\"k3\", \"k2\"]"});
  assertEqualVectors(expected, result);
}

} // namespace
} // namespace kumo::pollux::functions
